package com.pwser.utils;

import com.alibaba.fastjson.JSON;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

@Component
public class ESUtil {
    @Autowired
    private RestHighLevelClient elasticsearchRestHighLevelClient;
    private static RestHighLevelClient esRestClient;

    public static final String INDEX_ALIAS = "jd_alias";
    public static final String SORT = "";
    public static final int PAGE = 1;
    public static final int SIZE = 20;
    public static final String HIGHLIGHT_PRE_TAGS = "<span style='color:red'>";
    public static final String HIGHLIGHT_POST_TAGS = "</span>";
    public static final int QUERY_TIMEOUT = 60;

    @PostConstruct
    public void init() {
        esRestClient = elasticsearchRestHighLevelClient;
    }


    /**
     * 批量添加文档
     * @param list
     * @return
     * @throws IOException
     */
    public static Boolean bulkAddDocuments(List<?> list) throws IOException {
        BulkRequest bulkRequest = new BulkRequest();
        bulkRequest.timeout(new TimeValue(QUERY_TIMEOUT, TimeUnit.SECONDS));
        for (Object obj : list) {
            IndexRequest indexRequest = new IndexRequest(INDEX_ALIAS);
            indexRequest.source(JSON.toJSONString(obj), XContentType.JSON);
            bulkRequest.add(indexRequest);
        }
        BulkResponse bulkResponse = esRestClient.bulk(bulkRequest, RequestOptions.DEFAULT);
        return !bulkResponse.hasFailures();
    }

    /**
     * 查询所有数据
     * @param page
     * @return
     * @throws IOException
     */
    public static List<Map<String, Object>> matchAll(int page) throws IOException {
        QueryBuilder queryBuilder = QueryBuilders.matchAllQuery();
        List<Map<String, Object>> list = ESUtil.getResultBySearchRequest(queryBuilder, SORT, page, SIZE, null);
        return list;
    }

    /**
     * 关键字查询数据
     * @param field
     * @param keyword
     * @param page
     * @return
     * @throws IOException
     */
    public static List<Map<String, Object>> match(String field, String keyword, int page) throws IOException {
        QueryBuilder queryBuilder = QueryBuilders.matchQuery(field, keyword);
        HighlightBuilder highlightBuilder = ESUtil.createHighlightBuilder(field);
        List<Map<String, Object>> list = ESUtil.getResultBySearchRequest(queryBuilder, SORT, page, SIZE, highlightBuilder);
        return list;
    }

    /**
     * 公用查询请求获取结果
     * @param queryBuilder
     * @param sort
     * @param page
     * @param size
     * @param highlightBuilder
     * @return
     * @throws IOException
     */
    private static List<Map<String, Object>> getResultBySearchRequest(QueryBuilder queryBuilder,
                                                                      String sort,
                                                                      int page,
                                                                      int size,
                                                                      HighlightBuilder highlightBuilder) throws IOException {
        SearchRequest searchRequest = new SearchRequest(INDEX_ALIAS);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        if (queryBuilder != null) {
            sourceBuilder.query(queryBuilder);
        }
        if (sort != null && !"".equals(sort)) {
            String[] sortArr = sort.split(" ");
            sourceBuilder.sort(sortArr[0], sortArr[1].equals("desc") ? SortOrder.DESC : SortOrder.ASC);
        }
        if (page < 1) {
            page = 1;
        }
        sourceBuilder.from((page - 1) * size);
        sourceBuilder.size(size > 0 ? size : SIZE);
        if (highlightBuilder != null) {
            sourceBuilder.highlighter(highlightBuilder);
        }
        sourceBuilder.timeout(new TimeValue(QUERY_TIMEOUT, TimeUnit.SECONDS));
        searchRequest.source(sourceBuilder);
        SearchResponse searchResponse = esRestClient.search(searchRequest, RequestOptions.DEFAULT);
        List<Map<String, Object>> list = new ArrayList<>();
        for (SearchHit hit : searchResponse.getHits().getHits()) {
            HighlightField highlightField = hit.getHighlightFields().get("title");
            Map<String, Object> map = hit.getSourceAsMap();
            if (highlightField != null) {
                String highlight = "";
                for (Text text : highlightField.getFragments()) {
                    highlight += text.toString();
                }
                map.put("highlight", highlight);
            }
            list.add(map);
        }
        return list;
    }

    /**
     * 构建高亮
     * @param field
     * @return
     */
    private static HighlightBuilder createHighlightBuilder(String field) {
        HighlightBuilder highlight = new HighlightBuilder();
        highlight.preTags(HIGHLIGHT_PRE_TAGS);
        highlight.postTags(HIGHLIGHT_POST_TAGS);
        highlight.requireFieldMatch(false); // 关键字全部高亮
        highlight.field(field);
        return highlight;
    }

}
